import numpy as np
import matplotlib.pyplot as plt
from itertools import cycle
from sklearn.cluster import KMeans
from scipy.cluster.hierarchy import linkage, dendrogram, fcluster
from pylab import mpl
import os

# 设置显示中文字体
mpl.rcParams["font.sans-serif"] = ["SimHei"]
method_name = {
    "single": "最短距离法",
    "complete": "最长距离法",
    "average": "类平均法",
    "centroid": "重心法",
    "ward": "Ward法",
}


# 定义系统聚类法类
class HierarchicalClustering:
    def __init__(self, method="single", metric="euclidean"):
        """
        'euclidean'：欧氏距离，即样本之间的直线距离。
        'manhattan'：曼哈顿距离，即样本之间的城市街区距离（各坐标轴的绝对差值之和）。
        'chebyshev'：切比雪夫距离，即样本之间的各坐标轴绝对差值的最大值。
        'minkowski'：闵可夫斯基距离，是欧氏距离和曼哈顿距离的一种广义形式，可通过额外 p 参数指定.
        """
        self.metric = metric
        self.method_name = {
            "single": "最短距离法",
            "complete": "最长距离法",
            "average": "类平均法",
            "centroid": "重心法",
            "ward": "Ward法",
        }
        self.method = method

    def fit(self, data):
        self.data = data
        # 创建 linkage matrix
        self.Z = linkage(np.array(data), method=self.method, metric=self.metric)
        print(self.Z)

    def analysis(self, data, save=False, labels=False):
        # 通用的分析手段
        self.fit(data)
        self.plot_dendrogram(save, labels)

    def analysis_2d(self, data, k, save=False):
        # 若数据为2维，则可进一步可视化分析
        self.fit(data)
        if self.data.shape[1] != 2:
            raise ValueError("Data must be two-dimensional to plot clusters.")
        # 将层次聚类结果划分为 k 类
        labels = fcluster(self.Z, k, criterion="maxclust")
        # 绘制二维平面图
        colors = cycle("bgrcmyk")
        plt.figure(figsize=(8, 6))
        for i, color in zip(range(1, k + 1), colors):
            plt.scatter(
                self.data[labels == i, 0],
                self.data[labels == i, 1],
                c=color,
                label="Cluster %d" % i,
            )
        plt.xlabel("X")
        plt.ylabel("Y")
        plt.title(
            "Hierarchical Clustering(k=%d,%s)" % (k, self.method_name[self.method])
        )
        plt.legend(bbox_to_anchor=(1.0001, 1), loc="upper left")
        plt.subplots_adjust(right=0.8)
        if save:
            save_dir = "./figs/"
            os.makedirs(save_dir, exist_ok=True)
            save_path = os.path.join(
                save_dir,
                "2d_Hierarchical_Clustering_Dendrogram_Method_%s_k_%d.png"
                % (self.method_name[self.method], k),
            )
            plt.savefig(save_path)
            print("二维系统聚类图已保存", save_path)
        plt.show()

    def plot_dendrogram(self, save, labels):
        # 绘制谱系聚类图
        plt.figure(figsize=(10, 5))
        if labels:
            dendrogram(self.Z, labels=labels, orientation="right", show_contracted=True)
        else:
            dendrogram(
                self.Z,
                labels=[i + 1 for i in range(len(self.data))],
                orientation="right",
                show_contracted=True,
            )
        plt.ylabel("Data Index")
        plt.xlabel("Distance")
        plt.title(
            "Hierarchical Clustering Dendrogram(Method:%s)"
            % self.method_name[self.method]
        )
        if save:
            save_dir = "./figs/"
            os.makedirs(save_dir, exist_ok=True)
            save_path = os.path.join(
                save_dir,
                "Hierarchical_Clustering_Dendrogram_Method_%s.png"
                % self.method_name[self.method],
            )
            plt.savefig(save_path)
            print("谱系聚类图已保存", save_path)
        plt.show()

    def mul_save(data, save=True):
        # 用于批量保存图片
        for method in method_name.keys():
            hierarchical_cluster = HierarchicalClustering(method=method)
            hierarchical_cluster.analysis(data, save)

    def mul_save_2d(data, k, save=True):
        # 用于批量保存图片
        for method in method_name.keys():
            hierarchical_cluster = HierarchicalClustering(method=method)
            hierarchical_cluster.analysis_2d(data, k, save)


# 定义k-均值聚类法
class KMeans:
    def __init__(self, k, max_iters=100):
        self.k = k
        self.max_iters = max_iters

    def fit(self, data):
        # 先进行初始化，选择初始凝聚点
        self.centroids = self.initialize_centroids(data)
        # 进行循环，到最大迭代次数或者收敛了停止
        for _ in range(self.max_iters):
            # 计算每个样本到各个凝聚点的距离，选择最近的凝聚点
            clusters = self.assign_clusters(data)
            # 更新凝聚点
            new_centroids = self.update_centroids(data, clusters)
            # 检查是否收敛
            if np.allclose(new_centroids, self.centroids):
                break
            self.centroids = new_centroids
        print(clusters)
        return clusters

    def initialize_centroids(self, data):
        # 抽取初始凝聚点的样本标号(replace=False表示不放回抽样,没有重复)
        indices = np.random.choice(len(data), self.k, replace=False)
        return data[indices]

    def assign_clusters(self, data):
        """
        distances是一个矩阵[d_ij](n*k,n个样本聚k类),d_ij=||x_i-u_j||,x_i是第i个样本,u_j是第j个凝聚点.
        具体操作机制用到了广播,(5,2)->(5,1,2)->(5,2,2),然后对第二个维度求范数,得到(5,2)的矩阵.
        """
        distances = np.linalg.norm(data[:, np.newaxis] - self.centroids, axis=2)
        return np.argmin(distances, axis=1)

    def update_centroids(self, data, clusters):
        new_centroids = np.zeros_like(self.centroids)
        for i in range(self.k):
            new_centroids[i] = np.mean(data[clusters == i], axis=0)
        return new_centroids

    def plot_clusters(self, data, clusters, save=False):
        if data.shape[1] != 2:
            raise ValueError("Data must be two-dimensional to plot clusters.")
        colors = ["r", "g", "b", "y", "c", "m", "k"]  # 添加更多颜色以防不够用
        if self.k > len(colors):
            raise ValueError("Too many clusters for the number of available colors.")

        plt.figure(figsize=(8, 6))
        for i in range(self.k):
            plt.scatter(
                data[clusters == i][:, 0],
                data[clusters == i][:, 1],
                c=colors[i],
                label=f"Cluster {i+1}",
            )
        plt.scatter(
            self.centroids[:, 0],
            self.centroids[:, 1],
            marker="*",
            s=25,
            c="k",
            label="Centroids",
        )
        plt.xlabel("X")
        plt.ylabel("Y")
        plt.title("K-Means Clustering(k=%d)" % self.k)
        plt.legend(bbox_to_anchor=(1.0001, 1), loc="upper left")
        plt.subplots_adjust(right=0.8)
        if save:
            save_dir = "./figs/"
            os.makedirs(save_dir, exist_ok=True)
            save_path = os.path.join(save_dir, "KMeans_Dendrogram_k_%d.png" % self.k)
            plt.savefig(save_path)
            print("k均值聚类图已保存", save_path)
        plt.show()


def analysis_all_in_one_2d(data, k, save=False):
    if data.shape[1] != 2:
        raise ValueError("Data must be two-dimensional to plot clusters.")
    HierarchicalClustering.mul_save_2d(data, k, save)
    kmeans = KMeans(k)
    clusters = kmeans.fit(data)
    kmeans.plot_clusters(data, clusters, save)


if __name__ == "__main__":
    # 创建数据(labels可选)-----------------------------------------------------------------------

    data = np.array(
        [
            np.array([0.05798, 5.5150, 347.10, 21.910, 8586, 1742, 61.69]),
            np.array([0.08441, 3.9700, 347.20, 19.710, 7947, 2000, 2440]),
            np.array([0.07217, 1.1530, 54.85, 3.052, 3860, 1445, 9497]),
            np.array([0.15010, 1.7020, 307.50, 15.030, 12290, 1461, 6380]),
            np.array([5.74400, 2.8540, 229.60, 9.657, 8099, 266, 12520]),
            np.array([0.21300, 0.7058, 240.30, 13.910, 8980, 2820, 4135]),
        ]
    )
    data = np.array(
        [
            [109, 72],
            [88, 116],
            [54, 95],
            [192, 55],
            [202, 84],
            [216, 123],
            [131, 104],
            [55, 96],
            [142, 189],
            [236, 50],
            [176, 116],
            [211, 160],
            [243, 181],
            [272, 116],
            [282, 59],
            [346, 131],
            [286, 188],
            [347, 209],
            [391, 233],
            [321, 260],
            [470, 299],
            [580, 575],
            [603, 478],
            [548, 346],
            [571, 425],
            [624, 360],
            [681, 419],
            [704, 397],
            [700, 504],
            [825, 404],
            [790, 475],
            [564, 550],
            [542, 483],
            [575, 370],
            [540, 443],
            [684, 344],
            [711, 448],
            [621, 420],
            [658, 471],
            [791, 396],
            [750, 505],
            [592, 333],
        ]
    )
    # labels = ['Category A', 'Category B', 'Category C', 'Category D', 'Category E','Category F']

    # 创建系统聚类法对象-------------------------------------------------------------------------

    hierarchical_cluster = HierarchicalClustering(method="single")
    # method参数可选值:
    method_name = {
        "single": "最短距离法",
        "complete": "最长距离法",
        "average": "类平均法",
        "centroid": "重心法",
        "ward": "Ward法",
    }

    # 进行系统聚类分析并输出谱系聚类图--------------------------------------------------------------

    hierarchical_cluster.analysis(data)
    hierarchical_cluster.analysis_2d(data, 3)
    # 定义K均值聚类器并拟合数据

    kmeans = KMeans(k=3)
    clusters = kmeans.fit(data)

    # 绘制聚类结果
    kmeans.plot_clusters(data, clusters)
